import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA

# 1. 生成过去 100 天的销量数据
np.random.seed(42)  # 设置随机种子，保证结果可复现
days = pd.date_range(start='2023-01-01', periods=100, freq='D')
sales = np.random.randint(50, 150, size=100)  # 随机生成销量数据

# 2. 将数据存放在 DataFrame 中
data = pd.DataFrame({'date': days, 'sales': sales})
data.set_index('date', inplace=True)

# 3. 使用 ARIMA 模型进行预测
model = ARIMA(data['sales'], order=(5, 1, 0))  # order=(p,d,q)
model_fit = model.fit()

# 预测未来 10 天的销量
forecast_steps = 10
forecast = model_fit.forecast(steps=forecast_steps)

# 4. 将预测结果存储到 DataFrame
forecast_dates = pd.date_range(start=days[-1] + pd.Timedelta(days=1), periods=forecast_steps, freq='D')
forecast_data = pd.DataFrame({'date': forecast_dates, 'forecast_sales': forecast})
forecast_data.set_index('date', inplace=True)

# 5. 可视化结果
plt.figure(figsize=(10, 6))

# 绘制过去的销量数据
plt.plot(data.index, data['sales'], label='Past Sales')

# 绘制预测的销量数据
plt.plot(forecast_data.index, forecast_data['forecast_sales'], label='Forecast Sales', linestyle='--', color='orange')

# 标题和标签
plt.title('Sales Prediction Over Time')
plt.xlabel('Date')
plt.ylabel('Sales')
plt.legend()

# 展示图像
plt.grid(True)

# 保存图像到文件
plt.savefig('sales_prediction.png')

# 6. 保存数据
data.to_csv('past_sales.csv')  # 保存过去的销量数据
forecast_data.to_csv('forecast_sales.csv')  # 保存预测的销量数据
